package com.ppch.wuwamanus.advisors;

import org.springframework.ai.chat.client.advisor.api.*;
import org.springframework.ai.chat.prompt.PromptTemplate;
import reactor.core.publisher.Flux;

import java.util.HashMap;
import java.util.Map;

/**
 * Package:ppch-ai-agent
 * ClassName:ReReadingAdvisor
 *
 * @Author 泡泡茶壶
 * @Create 2025/7/3 20:54
 * @Version 1.0
 * Description:
 * 自定义重复阅读 Re2 顾问（拦截器）
 *   通过让大模型二次阅读用户问题，提高大型语言模型的推理能力
 */
public class ReReadingAdvisor implements CallAroundAdvisor, StreamAroundAdvisor {

    /**
     * 提供唯一的 advisor 名称。
     */
    @Override
    public String getName() {
        return this.getClass().getSimpleName();
    }

    /**
     * 确定链中的 advisor 顺序（值越小优先级越高）
     */
    @Override
    public int getOrder() {
        return 0;
    }

    /**
     * 前置顾问：让大模型二次阅读用户问题
     * @param advisedRequest 用户 Prompt（用户问题）
     */
    private AdvisedRequest before(AdvisedRequest advisedRequest) {
        //1.创建包含变量的模板字符串
        String template = """
        {re2_input_query}
        Read the question again: {re2_input_query}
        """;
        //2.使用 PromptTemplate 来渲染模板字符串
        PromptTemplate promptTemplate = new PromptTemplate(template);
        Map<String, Object> variables = new HashMap<>();
        variables.put("re2_input_query", advisedRequest.userText());
        String renderedText = promptTemplate.render(variables);
        return AdvisedRequest.from(advisedRequest)
                .userText(renderedText) // 使用渲染后的文本
                .userParams(variables)
                .build();
    }

    @Override
    public AdvisedResponse aroundCall(AdvisedRequest advisedRequest, CallAroundAdvisorChain chain) {
        return chain.nextAroundCall(this.before(advisedRequest));
    }

    @Override
    public Flux<AdvisedResponse> aroundStream(AdvisedRequest advisedRequest, StreamAroundAdvisorChain chain) {
        return chain.nextAroundStream(this.before(advisedRequest));
    }
}
